{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8900da4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/data/run01/sczc619/LML/MetaTSNE')\n",
    "from hogrl_main import *\n",
    "from hogrl_model import *\n",
    "from hogrl_utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "61ee97c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data from: /data/run01/sczc619/LML/MetaTSNE/data/Amazon_shortest_distance.pkl\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pickle\n",
    "\n",
    "# dataset_name = \"YelpChi\"\n",
    "dataset_name = \"Amazon\"\n",
    "\n",
    "def load_mindist_data(dataset_name):\n",
    "    current_dir = os.path.dirname(os.path.abspath(''))\n",
    "    data_path = os.path.join(current_dir, '..', 'data', f\"{dataset_name}_shortest_distance.pkl\")\n",
    "    data_path = os.path.abspath(data_path)\n",
    "    print(\"Loading data from:\", data_path)\n",
    "    with open(data_path, \"rb\") as f:\n",
    "        data = pickle.load(f)\n",
    "    return data[\"dist_matrix\"], data[\"sparse_matrix\"], data[\"node_count\"]\n",
    "\n",
    "\n",
    "dist_matrix, sparse_matrix, n_nodes = load_mindist_data(dataset_name)\n",
    "# print(sparse_matrix[0])\n",
    "# print(n_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "12f349b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import SAGEConv, GCNConv\n",
    "import os\n",
    "\n",
    "class multi_HOGRL_Model(nn.Module):\n",
    "    def __init__(self, in_feat, out_feat, relation_nums=3, hidden=32, \n",
    "                 drop_rate=0.6, weight=1, num_layers=2, layers_tree=2,\n",
    "                 tsne_weight=0.1):\n",
    "        super(multi_HOGRL_Model, self).__init__()\n",
    "        self.relation_nums=relation_nums\n",
    "        self.drop_rate = drop_rate\n",
    "        self.weight = weight\n",
    "        self.layers_tree = layers_tree\n",
    "        for i in range(relation_nums):\n",
    "            setattr(self,'Layers'+str(i),Layer_AGG(in_feat,hidden,self.drop_rate,self.weight,num_layers,self.layers_tree))\n",
    "        self.linear=nn.Linear(hidden*relation_nums,out_feat)\n",
    "        self.tsne_weight = tsne_weight\n",
    "        \n",
    "    def forward(self, x, edge_index):\n",
    "        layer_outputs = []\n",
    "\n",
    "        for i in range(self.relation_nums):\n",
    "            layer_output = getattr(self, 'Layers' + str(i))(x, edge_index[i])\n",
    "            layer_outputs.append(layer_output)\n",
    "\n",
    "        x_temp = torch.cat(layer_outputs, dim=1)\n",
    "\n",
    "        x = self.linear(x_temp)\n",
    "        x = F.log_softmax(x, dim=1)\n",
    "            \n",
    "        return x, x_temp\n",
    "    \n",
    "class Layer_AGG(nn.Module):\n",
    "    def __init__(self, in_feat, out_feat, drop_rate=0.6, weight=1, num_layers=2, layers_tree=2):\n",
    "        super().__init__()\n",
    "        self.drop_rate = drop_rate\n",
    "        self.weight = weight\n",
    "        self.num_layers = num_layers\n",
    "        self.layers_tree = layers_tree\n",
    "        \n",
    "        # 主路径\n",
    "        self.convs = nn.ModuleList()\n",
    "        for i in range(num_layers):\n",
    "            in_channels = in_feat if i==0 else out_feat\n",
    "            self.convs.append(nn.Sequential(\n",
    "                SAGEConv(in_channels, out_feat),\n",
    "                nn.BatchNorm1d(out_feat\n",
    "            )))\n",
    "            \n",
    "        if in_feat != out_feat:\n",
    "            self.residual = nn.Linear(in_feat, out_feat)\n",
    "        else:\n",
    "            self.residual = nn.Identity()\n",
    "            \n",
    "        self.conv_tree = nn.ModuleList()\n",
    "        self.gating_networks = nn.ModuleList()\n",
    "        for _ in range(layers_tree):\n",
    "            self.conv_tree.append(nn.Sequential(\n",
    "                SAGEConv(in_feat, out_feat),\n",
    "                nn.BatchNorm1d(out_feat)\n",
    "            ))\n",
    "            self.gating_networks.append(nn.Linear(out_feat, 1))\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        residual = self.residual(x) \n",
    "\n",
    "        h = x\n",
    "        for conv in self.convs:\n",
    "            h = conv[0](h, edge_index[0])\n",
    "            h = conv[1](h) \n",
    "            h = F.relu(h)\n",
    "            h = F.dropout(h, p=self.drop_rate, training=self.training)\n",
    "        h_main = h + residual \n",
    "\n",
    "        h_tree = []\n",
    "        for i in range(self.layers_tree):\n",
    "            h = self.conv_tree[i][0](x, edge_index[1][i])\n",
    "            h = self.conv_tree[i][1](h) \n",
    "            h = F.relu(h)\n",
    "            h = F.dropout(h, p=self.drop_rate, training=self.training)\n",
    "            h_tree.append(h)\n",
    "        \n",
    "        gates = torch.stack([gate(h) for gate, h in zip(self.gating_networks, h_tree)], dim=-1)\n",
    "        alpha = F.softmax(gates, dim=-1)\n",
    "        h_tree = sum([h * alpha[:, :, i] for i, h in enumerate(h_tree)])\n",
    "        \n",
    "        return h_main + self.weight * h_tree \n",
    "    \n",
    "class Graphsage(nn.Module):\n",
    "    def __init__(self, in_feat,out_feat):\n",
    "        super(Graphsage, self).__init__()\n",
    "        self.conv1 = SAGEConv(in_feat, out_feat)\n",
    "        self.conv2 = SAGEConv(out_feat, out_feat)\n",
    "        # self.conv1 = GCNConv(in_feat, out_feat)\n",
    "        # self.conv2 = GCNConv(out_feat, out_feat)\n",
    "        self.linear = nn.Linear(out_feat,2)\n",
    "\n",
    "\n",
    "    def forward(self,x,edge_index):\n",
    "        x = self.conv1(x,edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = F.dropout(x, p=0.6, training=self.training)\n",
    "        x = self.conv2(x,edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = F.dropout(x, p=0.6, training=self.training)\n",
    "        x = self.linear(x)\n",
    "        x = F.log_softmax(x,dim=1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f2d70dde",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_tsne_loss(embedding, dist_matrix, perplexity=30.0, temperature=10, alpha=1.0, df=1.0, eps=1e-12):\n",
    "    N = embedding.shape[0]\n",
    "    device = embedding.device\n",
    "    \n",
    "    distances = dist_matrix.float().clone()\n",
    "    distances.fill_diagonal_(0)\n",
    "    \n",
    "    sigma = find_sigma(distances, perplexity).view(-1, 1)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        P = torch.exp(-distances**2 / (2 * sigma**2))\n",
    "        P.fill_diagonal_(0)\n",
    "        P = (P + P.T) / (2 * N)\n",
    "        P = P / (P.sum(dim=1, keepdim=True) + eps)\n",
    "        P = torch.clamp(P, min=eps)\n",
    "\n",
    "    pairwise_dist = torch.cdist(embedding, embedding)\n",
    "    Q = (1 + pairwise_dist**2 / temperature)**-1 \n",
    "    Q.fill_diagonal_(0)\n",
    "    Q = (Q + Q.T) / (2 * N)\n",
    "    Q = Q / (Q.sum(dim=1, keepdim=True) + eps)\n",
    "    Q = torch.clamp(Q, min=eps)\n",
    "\n",
    "    loss = (P * torch.log(P / Q)).sum()\n",
    "    \n",
    "#     print(f\"Sigma mean: {sigma.mean().item()}, max: {sigma.max().item()}, min: {sigma.min().item()}\")\n",
    "#     print(f\"P mean: {P.mean().item()}, P max: {P.max().item()}, P min: {P.min().item()}\")\n",
    "#     print(f\"Q mean: {Q.mean().item()}, Q max: {Q.max().item()}, Q min: {Q.min().item()}\")\n",
    "    return loss\n",
    "    \n",
    "def find_sigma(distances, target_perplexity, tol=1e-5, max_iter=100):\n",
    "    distances = distances.detach()\n",
    "    sigmas = []\n",
    "    for i, row in enumerate(distances):\n",
    "        mask = row != 0\n",
    "        valid_distances = row[mask]\n",
    "        if len(valid_distances) == 0:\n",
    "            sigmas.append(0.0)\n",
    "            continue\n",
    "        \n",
    "        sigma_min, sigma_max = 1e-5, 10.0\n",
    "        best_sigma = sigma_max\n",
    "        for _ in range(max_iter):\n",
    "            sigma = (sigma_min + sigma_max) / 2\n",
    "            p = torch.exp(-valid_distances**2 / (2 * sigma**2))\n",
    "            p = p / (p.sum() + 1e-8)\n",
    "            entropy = -torch.sum(p * torch.log2(p + 1e-12))\n",
    "            current_perplexity = 2**entropy\n",
    "            \n",
    "            if abs(current_perplexity - target_perplexity) < tol:\n",
    "                best_sigma = sigma\n",
    "                break\n",
    "                \n",
    "            if current_perplexity > target_perplexity:\n",
    "                sigma_max = sigma\n",
    "            else:\n",
    "                sigma_min = sigma\n",
    "        sigmas.append(best_sigma)\n",
    "    return torch.tensor(sigmas, device=distances.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2ff5c086",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = {\n",
    "    \"dataset\": \"amazon\",\n",
    "    \"batch_size\": 256,\n",
    "    \"weight_decay\": 0.00005,\n",
    "    \"emb_size\": 64,\n",
    "    \"pretrain_epochs\": 200,   \n",
    "    \"pretrain_lr\": 0.005,      \n",
    "    \"finetune_lr\": 0.001,\n",
    "    \"num_epochs\": 500,\n",
    "    \"pretrain_patience\": 20,\n",
    "    \"patience\": 30,\n",
    "    \"tsne_weight\": 0.3,       \n",
    "    \"weight\": 0.6,\n",
    "    \"layers\": 7,\n",
    "    \"test_size\": 0.6,\n",
    "    \"val_size\": 0.5,\n",
    "    \"layers_tree\": 7,\n",
    "    \"seed\": 76,\n",
    "    \"drop_rate\": 0.5\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e09aeeac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n",
      "loading data...\n",
      "training...\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(device)\n",
    "print('loading data...')\n",
    "prefix = \"/data/run01/sczc619/LML/MetaTSNE/data/\"\n",
    "edge_indexs,feat_data,labels = load_data(args['dataset'],args['layers_tree'], prefix)\n",
    "np.random.seed(args['seed'])\n",
    "random.seed(args['seed'])\n",
    "\n",
    "if args['dataset'] == 'yelp':\n",
    "    index = list(range(len(labels)))\n",
    "    idx_train_val, idx_test, y_train_val, y_test = train_test_split(index, labels, stratify=labels, test_size=args['test_size'], random_state=2, shuffle=True)\n",
    "    idx_train, idx_val, y_train, y_val = train_test_split(idx_train_val, y_train_val, stratify=y_train_val, test_size=args['val_size'], random_state=2, shuffle=True)\n",
    "    dist_path = os.path.join(prefix, \"Yelp_shortest_distance.pkl\")\n",
    "elif args['dataset'] == 'amazon':\n",
    "    # 0-3304 are unlabeled nodes\n",
    "    index = list(range(3305, len(labels)))\n",
    "    idx_train_val, idx_test, y_train_val, y_test = train_test_split(index, labels[3305:], stratify=labels[3305:], test_size=args['test_size'], random_state=2, shuffle=True)\n",
    "    idx_train, idx_val, y_train, y_val = train_test_split(idx_train_val, y_train_val, stratify=y_train_val, test_size=args['val_size'], random_state=2, shuffle=True)\n",
    "    dist_path = os.path.join(prefix, \"Amazon_shortest_distance.pkl\")\n",
    "\n",
    "train_pos, train_neg = pos_neg_split(idx_train, y_train)\n",
    "\n",
    "# 加载距离矩阵\n",
    "#     dist_path = os.path.join(prefix, f\"{args['dataset']}_shortest_distance.pkl\")\n",
    "with open(dist_path, 'rb') as f:\n",
    "    dist_data = pickle.load(f)\n",
    "    dist_matrix = torch.tensor(dist_data['dist_matrix']).to(device)\n",
    "#         print(dist_matrix[0])\n",
    "\n",
    "#     dist_matrix = torch.tensor(dist_matrix).to(device)\n",
    "# 模型初始化\n",
    "gnn_model = multi_HOGRL_Model(\n",
    "    feat_data.shape[1],\n",
    "    2,\n",
    "    len(edge_indexs),\n",
    "    args['emb_size'],\n",
    "    args['drop_rate'],\n",
    "    args['weight'],\n",
    "    args['layers'],\n",
    "    args['layers_tree'], \n",
    "    tsne_weight=0.3 \n",
    ").to(device)\n",
    "\n",
    "for edge_index in edge_indexs:\n",
    "    edge_index[0] = edge_index[0].to(device)\n",
    "    edge_index[1] = [tensor.to(device) for tensor in edge_index[1]]\n",
    "\n",
    "# labels = torch.tensor(labels).to(device)\n",
    "feat_data = torch.tensor(feat_data).float().to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(gnn_model.parameters(), lr=0.005, weight_decay=5e-5)\n",
    "batch_size = args['batch_size']\n",
    "\n",
    "best_val_auc = 0.0\n",
    "best_model_state = None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d34449ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(args['num_epochs']):\n",
    "    gnn_model.train()\n",
    "#         total_loss = 0\n",
    "    tsne_loss = 0\n",
    "\n",
    "    # randomly under-sampling negative nodes for each epoch\n",
    "    sampled_idx_train = undersample(train_pos, train_neg, scale=1)\n",
    "    rd.shuffle(sampled_idx_train)\n",
    "\n",
    "    num_batches = int(len(sampled_idx_train) / batch_size ) + 1\n",
    "#         print(\"num_batches:\",num_batches)\n",
    "\n",
    "    for batch in range(num_batches):\n",
    "        i_start = batch * batch_size\n",
    "        i_end = min((batch + 1) * batch_size, len(sampled_idx_train))\n",
    "        batch_nodes = sampled_idx_train[i_start:i_end]\n",
    "        batch_label = torch.tensor(labels[np.array(batch_nodes)]).long().to(device)\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a97bec30",
   "metadata": {},
   "outputs": [],
   "source": [
    "out, f = gnn_model(feat_data, edge_indexs) # out 是预测，f是产生的特征"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89ac2e29",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ad951de1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(256, 256)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(batch_nodes), batch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab60742c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19b42674",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7f21e4c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 2, 2,  ..., 2, 2, 2],\n",
       "        [2, 0, 2,  ..., 2, 2, 2],\n",
       "        [2, 2, 0,  ..., 2, 2, 2],\n",
       "        ...,\n",
       "        [2, 2, 2,  ..., 0, 1, 2],\n",
       "        [2, 2, 2,  ..., 1, 0, 2],\n",
       "        [2, 2, 2,  ..., 2, 2, 0]], device='cuda:0', dtype=torch.int32)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dist_matrix[batch_nodes][:, batch_nodes]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "849a9e01",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([11944, 2])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4ee50861",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([11944, 192])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a286a2c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tsne_loss 单独测试\n",
    "def hogrl_main(args):\n",
    "\n",
    "    \n",
    "\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "#             out,_ = gnn_model(feat_data,edge_indexs)\n",
    "            \n",
    "            f.retain_grad()\n",
    "        \n",
    "#             print(f\"f norm: {torch.norm(f).item()}\")\n",
    "            \n",
    "            # 只保留tsne_loss\n",
    "#             tsne_loss = torch.tensor(0.0).to(out.device)\n",
    "            tsne_loss = calculate_tsne_loss(f, dist_matrix) * gnn_model.tsne_weight\n",
    "            \n",
    "#             print(tsne_loss)\n",
    "            batch_nodes_tensor = torch.tensor(batch_nodes, dtype=torch.long, device=device)\n",
    "\n",
    "            tsne_loss.backward()\n",
    "\n",
    "            print(\"tsne_loss:\",tsne_loss)\n",
    "            print(f\"Embedding grad stats: mean={f.grad.mean():.3e}, std={f.grad.std():.3e}\")\n",
    "            \n",
    "            optimizer.step()\n",
    "            \n",
    "        \n",
    "#             tsne_loss = torch.tensor(0.0).to(out.device)\n",
    "#             tsne_loss = calculate_tsne_loss(f, dist_matrix) * gnn_model.tsne_weight\n",
    "            \n",
    "# #             print(tsne_loss)\n",
    "#             batch_nodes_tensor = torch.tensor(batch_nodes, dtype=torch.long, device=device)\n",
    "\n",
    "#             # loss = F.nll_loss(out[np.array(batch_nodes)], batch_label)\n",
    "#             cls_loss = F.nll_loss(out[batch_nodes_tensor], batch_label)\n",
    "#             total_loss = cls_loss + tsne_loss\n",
    "#             total_loss = total_loss.sum()\n",
    "# #             print(total_loss)\n",
    "#             total_loss.backward()\n",
    "#             if f.grad is not None:\n",
    "#                 print(f\"Embedding grad norm: {torch.norm(f.grad).item():.4f}\")\n",
    "#             else:\n",
    "#                 print(\"Embedding f gradient is None!\")\n",
    "#             print(\"tsne_loss:\",tsne_loss)\n",
    "# #             print(f\"tsne_loss: {tsne_loss.item():.4f}, total_loss: {total_loss.item():.4f}\")\n",
    "            \n",
    "# #             loss.backward()\n",
    "#             optimizer.step()\n",
    "#             total_loss += total_loss.item()\n",
    "        \n",
    "#             loss += loss.item()\n",
    "        if epoch % 10 == 9: # validate every 10 epochs \n",
    "            val_auc, val_ap, val_f1, val_g_mean = test(idx_val, y_val, gnn_model, feat_data, edge_indexs)\n",
    "            print(f'Epoch: {epoch}, Val AUC: {val_auc:.4f}, Val AP: {val_ap:.4f}, Val F1: {val_f1:.4f}, Val G-mean: {val_g_mean:.4f}')\n",
    "            \n",
    "            if val_auc > best_val_auc:\n",
    "                best_val_auc = val_auc\n",
    "                best_model_state = gnn_model.state_dict() \n",
    "                \n",
    "    gnn_model.load_state_dict(best_model_state)  \n",
    "    test_auc, test_ap, test_f1, test_g_mean = test(idx_test, y_test, gnn_model, feat_data, edge_indexs)\n",
    "    print(f'Test AUC: {test_auc:.4f}, Test AP: {test_ap:.4f}, '\n",
    "        f'Test F1: {test_f1:.4f}, Test G-mean: {test_g_mean:.4f}')\n",
    "    out,embedding = gnn_model(feat_data,edge_indexs)\n",
    "\n",
    "args = {\n",
    "    \"dataset\": \"amazon\",\n",
    "    \"batch_size\": 256,\n",
    "    \"lr\": 0.005,\n",
    "    \"weight_decay\": 0.00005,\n",
    "    \"emb_size\": 64,\n",
    "    \"num_epochs\": 500,\n",
    "    \"weight\": 0.6,\n",
    "    \"layers\": 7,\n",
    "    \"test_size\": 0.6,\n",
    "    \"val_size\": 0.5,\n",
    "    \"layers_tree\": 7,\n",
    "    \"seed\": 76,\n",
    "    \"drop_rate\": 0.3\n",
    "}\n",
    "\n",
    "hogrl_main(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34f784e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 自监督预训练 + 任务损失微调 \n",
    "\n",
    "import time\n",
    "import copy\n",
    "import random\n",
    "import torch\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import numpy as np\n",
    "import random as rd\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import f1_score, accuracy_score, recall_score, roc_auc_score, average_precision_score\n",
    "\n",
    "def hogrl_main(args):\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    print(device)\n",
    "    \n",
    "    timestamp = time.strftime(\"%Y%m%d-%H%M%S\")\n",
    "    writer = SummaryWriter(f'runs/{args[\"dataset\"]}_{timestamp}')\n",
    "    \n",
    "    print('loading data...')\n",
    "    prefix = \"/data/run01/sczc619/LML/MetaTSNE/data/\"\n",
    "    edge_indexs,feat_data,labels = load_data(args['dataset'],args['layers_tree'], prefix)\n",
    "                                             \n",
    "    np.random.seed(args['seed'])\n",
    "    random.seed(args['seed'])\n",
    "    \n",
    "    if args['dataset'] == 'yelp':\n",
    "        index = list(range(len(labels)))\n",
    "        idx_train_val, idx_test, y_train_val, y_test = train_test_split(index, labels, stratify=labels, \n",
    "                                                                       test_size=args['test_size'], random_state=2, shuffle=True)\n",
    "        idx_train, idx_val, y_train, y_val = train_test_split(idx_train_val, y_train_val, \n",
    "                                                             stratify=y_train_val, test_size=args['val_size'], \n",
    "                                                             random_state=2, shuffle=True)\n",
    "        dist_path = os.path.join(prefix, \"Yelp_shortest_distance.pkl\")\n",
    "    elif args['dataset'] == 'amazon':\n",
    "        index = list(range(3305, len(labels)))\n",
    "        idx_train_val, idx_test, y_train_val, y_test = train_test_split(index, labels[3305:], \n",
    "                                                                       stratify=labels[3305:], test_size=args['test_size'], \n",
    "                                                                       random_state=2, shuffle=True)\n",
    "        idx_train, idx_val, y_train, y_val = train_test_split(idx_train_val, y_train_val, \n",
    "                                                            stratify=y_train_val, test_size=args['val_size'], \n",
    "                                                            random_state=2, shuffle=True)\n",
    "        dist_path = os.path.join(prefix, \"Amazon_shortest_distance.pkl\")\n",
    "        \n",
    "    with open(dist_path, 'rb') as f:\n",
    "        dist_data = pickle.load(f)\n",
    "        dist_matrix = torch.tensor(dist_data['dist_matrix']).to(device)\n",
    "    \n",
    "    gnn_model = multi_HOGRL_Model(\n",
    "        feat_data.shape[1],\n",
    "        2,\n",
    "        len(edge_indexs),\n",
    "        args['emb_size'],\n",
    "        args['drop_rate'],\n",
    "        args['weight'],\n",
    "        args['layers'],\n",
    "        args['layers_tree'], \n",
    "        tsne_weight=args['tsne_weight']\n",
    "    ).to(device)\n",
    "    \n",
    "    for edge_index in edge_indexs:\n",
    "        edge_index[0] = edge_index[0].to(device)\n",
    "        edge_index[1] = [tensor.to(device) for tensor in edge_index[1]]\n",
    "    feat_data = torch.tensor(feat_data).float().to(device)\n",
    "\n",
    "    # 自监督预训练\n",
    "    print(\"\\n=== Starting Pretraining ===\")\n",
    "    gnn_model.linear.requires_grad_(False)\n",
    "#     optimizer = torch.optim.Adam(\n",
    "#         filter(lambda p: p.requires_grad, gnn_model.parameters()),\n",
    "#         lr=args['pretrain_lr'],\n",
    "#         weight_decay=5e-5\n",
    "#     )\n",
    "    optimizer = torch.optim.AdamW(  \n",
    "        filter(lambda p: p.requires_grad, gnn_model.parameters()),\n",
    "        lr=args['pretrain_lr'],\n",
    "        weight_decay=5e-5\n",
    "    )\n",
    "    pretrain_best_loss = float('inf')\n",
    "    pretrain_no_improve = 0\n",
    "    pretrain_early_stop = False\n",
    "    \n",
    "    \n",
    "    for epoch in range(args['pretrain_epochs']):\n",
    "        if pretrain_early_stop:\n",
    "            break\n",
    "            \n",
    "        gnn_model.train()\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        _, embeddings = gnn_model(feat_data, edge_indexs)\n",
    "        \n",
    "        tsne_loss = calculate_tsne_loss(embeddings, dist_matrix) * args['tsne_weight']\n",
    "        tsne_loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if tsne_loss.item() < pretrain_best_loss:\n",
    "            pretrain_best_loss = tsne_loss.item()\n",
    "            pretrain_no_improve = 0\n",
    "        else:\n",
    "            pretrain_no_improve += 1\n",
    "            \n",
    "        if pretrain_no_improve >= args['pretrain_patience']:\n",
    "            print(f\"Pretrain early stopping at epoch {epoch}\")\n",
    "            pretrain_early_stop = True\n",
    "        \n",
    "        writer.add_scalar('Pretrain/TSNE_Loss', tsne_loss.item(), epoch)\n",
    "        \n",
    "        if epoch % 5 == 0:\n",
    "            print(f'Pretrain Epoch: {epoch:03d}, TSNE Loss: {tsne_loss.item():.4f}')\n",
    "\n",
    "    # 微调（分类损失）\n",
    "    print(\"\\n=== Starting Fine-tuning ===\")\n",
    "    gnn_model.linear.requires_grad_(True)\n",
    "#     optimizer = torch.optim.Adam(\n",
    "#         gnn_model.parameters(),\n",
    "#         lr=args['finetune_lr'],\n",
    "#         weight_decay=5e-5\n",
    "#     )\n",
    "    optimizer = torch.optim.AdamW(\n",
    "        filter(lambda p: p.requires_grad, gnn_model.parameters()),\n",
    "        lr=args['pretrain_lr'],\n",
    "        weight_decay=5e-5\n",
    "    )\n",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "        optimizer, mode='max', factor=0.5, patience=10, verbose=True\n",
    "    )\n",
    "    \n",
    "    best_val_auc = 0.0\n",
    "    best_model_state = None\n",
    "    train_pos, train_neg = pos_neg_split(idx_train, y_train)\n",
    "    \n",
    "    no_improve_epochs = 0\n",
    "    early_stop = False\n",
    "    \n",
    "    for epoch in range(args['num_epochs']):\n",
    "        if early_stop:\n",
    "            break\n",
    "            \n",
    "        gnn_model.train()\n",
    "        total_loss = 0.0\n",
    "        \n",
    "        sampled_idx_train = undersample(train_pos, train_neg, scale=1)\n",
    "        rd.shuffle(sampled_idx_train)\n",
    "        \n",
    "        num_batches = int(len(sampled_idx_train) / args['batch_size']) + 1\n",
    "        \n",
    "        for batch in range(num_batches):\n",
    "            i_start = batch * args['batch_size']\n",
    "            i_end = min((batch + 1) * args['batch_size'], len(sampled_idx_train))\n",
    "            batch_nodes = sampled_idx_train[i_start:i_end]\n",
    "            batch_label = torch.tensor(labels[np.array(batch_nodes)]).long().to(device)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            out, _ = gnn_model(feat_data, edge_indexs) \n",
    "            \n",
    "            cls_loss = F.nll_loss(out[batch_nodes], batch_label)\n",
    "            \n",
    "            cls_loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            total_loss += cls_loss.item()\n",
    "\n",
    "        avg_loss = total_loss / num_batches\n",
    "        writer.add_scalar('FineTune/Train_Loss', avg_loss, epoch)\n",
    "        \n",
    "        if epoch % 5 == 0:\n",
    "            val_auc, val_ap, val_f1, val_g_mean = test(idx_val, y_val, gnn_model, feat_data, edge_indexs)\n",
    "            \n",
    "            writer.add_scalar('Validation/AUC', val_auc, epoch)\n",
    "            writer.add_scalar('Validation/F1', val_f1, epoch)\n",
    "            writer.add_scalar('Validation/GMean', val_g_mean, epoch)\n",
    "            \n",
    "            print(f'Epoch: {epoch:03d} | Loss: {avg_loss:.4f} | Val AUC: {val_auc:.4f} | Val F1: {val_f1:.4f}')\n",
    "            \n",
    "            scheduler.step(val_auc)\n",
    "            \n",
    "            if val_auc > best_val_auc:\n",
    "                best_val_auc = val_auc\n",
    "                no_improve_epochs = 0\n",
    "                best_model_state = copy.deepcopy(gnn_model.state_dict())\n",
    "            else:\n",
    "                no_improve_epochs += 5\n",
    "                \n",
    "            if no_improve_epochs >= args['patience']:\n",
    "                print(f\"Early stopping at epoch {epoch}\")\n",
    "                early_stop = True\n",
    "                \n",
    "                \n",
    "    gnn_model.load_state_dict(best_model_state)\n",
    "    test_auc, test_ap, test_f1, test_g_mean = test(idx_test, y_test, gnn_model, feat_data, edge_indexs)\n",
    "    print(f'\\n=== Final Test Results ===')\n",
    "    print(f'Test AUC: {test_auc:.4f} | Test AP: {test_ap:.4f} | Test F1: {test_f1:.4f} | G-mean: {test_g_mean:.4f}')\n",
    "    writer.close()\n",
    "                                             \n",
    "#     print('Generating embedding visualization...')\n",
    "#     _, final_embeddings = gnn_model(feat_data, edge_indexs)\n",
    "#     Visualization(labels, final_embeddings.cpu().detach(), prefix)\n",
    "                                             \n",
    "                                            \n",
    "\n",
    "\n",
    "hogrl_main(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7645d806",
   "metadata": {},
   "source": [
    "### Batch 处理\n",
    "\n",
    "基于KL散度的条件概率对齐：\n",
    "$$\n",
    "P(x, y) = \\frac{x}{}\n",
    "$$\n",
    "\n",
    "\n",
    "\n",
    "batch处理的理论依据：\n",
    "\n",
    "\n",
    "$$\n",
    "\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cd4ca6a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "230ca9e8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tsne",
   "language": "python",
   "name": "tsne"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
